## This code was kindly provided by Dan Corstange as part of the code
## that replicates Corstange, Daniel, "Sensitive Questions, Truthful Answers?
## Modeling the List Experiment with LISTIT," Political Analysis 17(1): 45-63.
## The code is identical to Corstange's replication code except line ##, noted.
## The change was made to report more parts of the regression output.

listit <- function(Y, X, map=NULL, method="BFGS") {
  ## Returns a matrix of parameters with a row for each list item
  ## and columns for each covariate.  The parameters are placed in
  ## their appropriate cells, and the remaining cells are filled
  ## with zeros.
  place.B <- function(param, map) {
    k <- nrow(map)
    n.x <- ncol(map)
    n.param <- length(param)
    B <- matrix(NA, nrow=k,
                    ncol=n.x)
    i.param <- 1
    for (i in 1:k)
      for (j in 1:n.x) {
        if (map[i,j] == TRUE) {
          B[i,j] <- param[i.param]
          i.param <- i.param + 1
        }
        else
          B[i,j] <- 0
      }
    return(B)
  }
  
  neg.LL <- function(param, Y.t, Y.c, X.t, X.c, map) {
    k <- nrow(map)
    B <- place.B(param=param, map=map)
    phi <- apply((1 + exp(-X.t %*%
                          t(B[2:k,])))^(-1)
                 , 1, sum)
    p <- as.vector(((1 + exp(-X.t %*%
                             B[1,]))^(-1) +
                    phi) / k)
    q <- (1 + exp(-X.c %*% t(B[2:k,])))^(-1)
    list.LL <- sum(log(choose(k, Y.t)) +
                   Y.t * log(p) +
                   (k - Y.t) * log(1 - p))
    off.LL <- sum(Y.c * log(q) +
                  (1 - Y.c) * log(1 - q))
    joint.LL <- list.LL + off.LL
    return(-joint.LL)
  }
  
  neg.grad <-function(param, Y.t, Y.c, X.t, X.c, map) {
    k <- nrow(map)                      #number of list items
    n.x <- ncol(map)                    #number of covariates
    n.grad <- sum(map)                  #number of gradients to return
    B <- place.B(param=param, map=map)  #rows: list item, cols: parameters
    phi <- apply((1 + exp(-X.t %*%
                          t(B[2:k,])))^(-1)
                 , 1, sum)
    p <- as.vector(((1 + exp(-X.t %*%
                             B[1,]))^(-1) +
                    phi) / k)
    ## Calculate the derivative of p with respect to the parameters
    ## Loop through the list items
    ## For each list item, get a temporary matrix of its predictors (this.X)
    ## Use (i.grad) and (this.param) indices to make sure the results go into
    ##   the appropriate columns in the (d.p) matrix, which has a number of
    ##   rows equal to the number of treatment observations, and a number of
    ##   columns equal to the number of gradients to be calculated (i.e., the
    ##   number of parameters we're trying to estimate).
    i.grad <- 0                         #gradient index
    d.p <- matrix(0, nrow=nrow(Y.t), ncol=n.grad)
    for (i in 1:k) {
      this.X <- X.t[,1]
      this.param <- sum(map[i,])
      for (j in 2:n.x)
        if (map[i,j] == TRUE)
          this.X <- cbind(this.X, X.t[,j]) #only the Xs for this list item
      e.vec <-
        as.vector(exp(-X.t %*% B[i,]) /
                  (k * (1 + exp(-X.t %*% B[i,]))^(2)))
      d.p[,(1 + i.grad):
            (i.grad + this.param)] <-
              this.X * e.vec
      i.grad <- i.grad + this.param
    }
    ## Calculate the contributions to the gradient calculations from the
    ## off items (i.e., the non-sensitive list items).
    ## Loop through the off items (items 2 to k).
    ## For each item, get a temporary matrix of its predictors (this.X).
    ## Use (i.grad) and (this.param) indices to make sure the results go into
    ##   the appropriate column in the (off) matrix.  (Note that the (off)    
    ##   matrix leaves a number of initial columns as zeros, corresponding to
    ##   the sensitive betas which do not enter these calculations.)  The
    ##   number of rows is the number of control group observations, and the
    ##   number of columns is the number of gradients to be calculated (i.e.,
    ##   the number of parameters we're trying to estimate).
    i.grad <- sum(map[1,])              #Start after the sensitive betas
    off <- matrix(0, nrow=nrow(Y.c), ncol=n.grad)
    for (i in 2:k) {
      this.X <- X.c[,1]
      this.param <- sum(map[i,])
      for (j in 2:n.x)
        if (map[i,j] == TRUE)
          this.X <- cbind(this.X, X.c[,j])
      q <- as.vector((1 + exp(-X.c %*%
                              B[i,]))^(-1))
      off[,(1 + i.grad):(i.grad + this.param)] <-
        (Y.c[,(i - 1)] - q) * this.X
      i.grad <- i.grad + this.param
    }
    ## Calculate all the relevant components and add them together.
    ## (LL.t) is the component related to the treatment group.
    ## (LL.off) is the component related to the off items.
    ## (LL.joint) is the joint components.
    LL.t <- as.vector(Y.t * p^(-1)) * d.p +
      as.vector((k - Y.t) * (1 - p)^(-1)) * (-d.p)
    LL.t <- apply(LL.t, 2, sum)
    LL.off <- apply(off, 2, sum)
    LL.joint <- LL.t + LL.off
    return(-LL.joint)
  }
  
  ## Removes observations that have missing values from the
  ## working data.
  na.rm <- function(Y=Y, X=X) {
    k <- ncol(Y)                          #number of list items
    x.names <- colnames(X)                #names of the covariates
    X <- cbind(1,X)                       #append a constant to X
    colnames(X) <- c("constant", x.names) #names the columns
    ## D is a data frame.  The first column will be a vector of
    ## dummies indicating whether these observations will be used
    ## or not depending on whether or not they have NAs among the
    ## covariates.  The second column marks whether the observation
    ## is in the treatment group (ie, it gets the list) or the
    ## control group (it gets the off items individually).  The
    ## Y values follow, then the X values (with an appended constant).
    D <- as.data.frame(cbind(NA,NA,Y,X))
    colnames(D) <- c("to.use", "control",
                     colnames(Y),
                     colnames(X))
    ## If one of the off items is NA, puts NA in control, else a number
    D[,"control"] <-
      apply(D[,(4:(k + 2))], 1, sum)
    ## If control is not NA, it's a control variable, else it's NA...
    D[,"control"] <-
      ifelse(!is.na(D[,"control"]),
             TRUE, D[,"control"])
    ## ...But if there is a list count, it's not a control variable
    D[,"control"] <-
      ifelse(!is.na(D[,3]),
             FALSE, D[,"control"])
    ## Puts a count of the number of covariate NAs into to.use
    D[,"to.use"] <-
      apply(is.na(D[,((k + 3):ncol(D))]),
            1, sum)
    ## Mark to.use if there are no covariate NAs on this obs
    D[,"to.use"] <-
      ifelse((D[,"to.use"] != 0),
             FALSE, TRUE)
    ## Mark to.use as FALSE if control observations have NA Y values
    D[,"to.use"] <-
      ifelse(is.na(D[,"control"]),
             FALSE, D[,"to.use"])
    ## Take the subset of usable observations, stripping off
    ## the to.use and control columns
    D <- subset(D[3:ncol(D)], (D[,"to.use"] == TRUE))
    return(D)                             #return the usable Y and X values
  }

  ## Checks for the map defining which covariates predict which list items.
  ## If no map is supplied by the user, it assumes that every covariate in X
  ## predicts every list item.  If a map is supplied by a user, it appends
  ## an initial column of "TRUE" values to indicate that the constant is
  ## a covariate for every list item.
  if (is.null(map))
    map <- matrix(TRUE, nrow=ncol(Y),
                  ncol=(ncol(X) + 1))
  else map <- cbind(TRUE, map)
  ## Processes the data.  First it removes observations with NA
  ## values, then places into separate objects control and treatment
  ## group observations on Y and X.
  k <- ncol(Y)
  YX <- na.rm(Y=Y, X=X)
  YX.c <- subset(YX, is.na(YX[,1]))
  YX.t <- subset(YX, !is.na(YX[,1]))
  Y.c <- as.matrix(YX.c[,2:k])
  X.c <- as.matrix(YX.c[(k + 1):ncol(YX.c)])
  Y.t <- as.matrix(YX.t[,1])
  X.t <- as.matrix(YX.t[(k + 1):ncol(YX.t)])
  ## The number of parameters are the sum total TRUE values
  ## in the parameter map
  n.param <- sum(map)
  
  result <- optim(par=rep(0, n.param),
                  fn=neg.LL, gr=neg.grad,
                  hessian=TRUE,
                  method=method,
                  Y.t=Y.t, Y.c=Y.c,
                  X.t=X.t, X.c=X.c,
                  map=map)

  neg.se <- FALSE
  for (i in 1:ncol(X.t))
    if (is.nan(sqrt(diag(solve(result$hessian)))[i])) {
      neg.se <- TRUE
      break
    }

  output <- list(b.star=(result$par)[1:ncol(X.t)],
                 var=solve(result$hessian),
                 se=sqrt(diag(solve(result$hessian)))[1:ncol(X.t)],
                 neg.se=neg.se,
                 converged=(result$convergence == 0),
                 n.treatment=nrow(Y.t),
                 n.control=nrow(Y.c),
                 items=k,
                 logL=-result$value,
                 deviance=2 * result$value,  ## ")" was removed and "," added to the end of this line
                 ## the following two lines were added:
                 b.control=(result$par)[(ncol(X.t)+1):length(result$par)],
                 se.control=sqrt(diag(solve(result$hessian)))[(ncol(X.t)+1):length(result$par)])

  class(output) <- "listit"
  
  return(output)
}
